import secrets
import hashlib
import os
from math_utils import (
    generate_prime_number,
    extended_gcd,
    mod_inverse,
    mgf1,
    bytes_xor,
    bytes_to_int,
    int_to_bytes,
    power
)

class RSA:
    """
    RSA 加密算法实现，支持 OAEP 填充 (PKCS#1 v2.2)。
    """

    def __init__(self, bits=2048, hash_func=hashlib.sha256):
        if bits < 16:
            raise ValueError("密钥位数太小，至少应为16位")
        self.hash_func = hash_func
        self.public_key, self.private_key = self._generate_keypair(bits)

    @staticmethod
    def _generate_keypair(bits):
        """
        生成 RSA 公私钥对。
        """
        p_bits = bits // 2
        q_bits = bits - p_bits
        p = generate_prime_number(p_bits)
        q = generate_prime_number(q_bits)
        while p == q:
            q = generate_prime_number(q_bits)
        n = p * q
        phi_n = (p - 1) * (q - 1)
        e = 65537
        # 确保 e 与 phi_n 互质
        while extended_gcd(e, phi_n)[0] != 1:
            # 修改: 使用 secrets 模块安全地选择备用公钥 e
            # 此代码路径极少被执行, 但仍需保证安全
            # 生成一个 [3, phi_n-1] 范围内的奇数
            e = 2 * secrets.randbelow((phi_n - 3) // 2) + 3
            
        d = mod_inverse(e, phi_n)
        return (n, e), (n, d)

    def encrypt(self, plaintext_bytes, label=b""):
        n, e = self.public_key
        k = (n.bit_length() + 7) // 8
        em = self._oaep_encode(plaintext_bytes, k, label)
        m_int = bytes_to_int(em)
        c = power(m_int, e, n)
        return c

    def decrypt(self, ciphertext_int, label=b""):
        n, d = self.private_key
        k = (n.bit_length() + 7) // 8
        if not (0 <= ciphertext_int < n):
            raise ValueError("密文不合法: 超出范围")
        m_int = power(ciphertext_int, d, n)
        em = int_to_bytes(m_int, min_length=k)
        return self._oaep_decode(em, k, label)

    def _oaep_encode(self, m, k, label):
        h = self.hash_func
        m_len = len(m)
        h_len = h().digest_size
        if m_len > k - 2*h_len - 2:
            raise ValueError("消息过长，无法进行 OAEP 填充")
        l_hash = h(label).digest()
        ps = b"\x00" * (k - m_len - 2*h_len - 2)
        db = l_hash + ps + b"\x01" + m
        # os.urandom 是安全的，无需修改
        seed = os.urandom(h_len)
        db_mask = mgf1(seed, k - h_len - 1, h)
        masked_db = bytes_xor(db, db_mask)
        seed_mask = mgf1(masked_db, h_len, h)
        masked_seed = bytes_xor(seed, seed_mask)
        em = b"\x00" + masked_seed + masked_db
        return em

    def _oaep_decode(self, em, k, label):
        h = self.hash_func
        h_len = h().digest_size
        if len(em) != k or k < 2*h_len + 2:
            raise ValueError("OAEP 解码失败: 长度不匹配")
        l_hash = h(label).digest()
        y = em[0:1]
        masked_seed = em[1:1+h_len]
        masked_db = em[1+h_len:]
        seed_mask = mgf1(masked_db, h_len, h)
        seed = bytes_xor(masked_seed, seed_mask)
        db_mask = mgf1(seed, k - h_len - 1, h)
        db = bytes_xor(masked_db, db_mask)
        if db[:h_len] != l_hash or y != b"\x00":
            raise ValueError("OAEP 解码失败: 哈希校验不通过")
        idx = db.find(b"\x01", h_len)
        if idx < 0:
            raise ValueError("OAEP 解码失败: 分隔符不存在")
        return db[idx+1:]